# -*- coding: utf-8 -*-
# SSHexec asyncssh工具文件

# 系统或第三方模块
import time
import signal
import asyncssh
import asyncio

# 自定义模块
import src.config as config
import src.utils as utils
from src.utils import tlog, elog

# 全局中断标志
asyncssh_interrupted = False 

async def connect_with_retry(runner, connect_timeout, max_retries=3):
    """
    功能：
        异步连接SSH服务器，支持重试机制

    参数：
        runner: 包含SSH连接信息的字典，必须包含'ip', 'port', 'user', 'password'
        connect_timeout: 连接超时时间，单位秒
        max_retries: 最大重试次数，默认3次

    返回：
        asyncssh.Connection: 成功连接后的SSH连接对象

    异常：
        ConnectionError: 所有重试都失败时抛出
    """
    if max_retries < 1:
        raise ValueError("max_retries must be at least 1")

    last_exception = None
    for attempt in range(max_retries):
        try:
            return await asyncssh.connect(
                host=runner["ip"],
                port=runner["port"],
                username=runner["user"],
                password=runner["password"],
                known_hosts=None,
                connect_timeout=connect_timeout
            )

        except (TimeoutError, asyncio.exceptions.TimeoutError) as e:
            elog.error(f"{runner.get('ip')}：连接超时，当前连接尝试次数 {attempt + 1}，将在 {2 ** attempt} 秒后重试")
            last_exception = e
            if attempt < max_retries - 1:
                await asyncio.sleep(2 ** attempt)  # 指数退避
        
        except Exception as e:
            # 其他错误，不是连接超时，退出重试
            elog.error(f"{runner.get('ip')}： 连接错误: {e}")
            last_exception = e
            raise last_exception
    
    # 确保 last_exception 不是 None
    if last_exception is None:
        raise RuntimeError("Unexpected error: No exception was captured")
    
    elog.error(f"所有重试都失败")
    raise last_exception  


def add_env_sudo_to_commands(args):
    """
    根据参数包装命令字符串，添加环境变量和sudo权限
    
    Args:
        args: 参数字典，包含：
            - c: 命令字符串
            - m: 模式，检测到'sudo'时添加sudo
            - e: 环境变量字符串
    
    Returns:
        str: 组合后的完整命令字符串
    """
    import shlex

    elog.info(f"开始拼接完整可执行命令")
    # 初始化解释器为None
    interpreter = None
    # 检查是否需要sudo
    sudo_prefix = 'sudo' if args.m == "sudo" else ''
    # 构建环境变量部分
    env_prefix = f"{args.e}" if args.e else ''
    # 构建超时部分
    timeout_prefix = f"timeout {args.t}" if args.t else ''
    
    if args.c:
        # 安全处理命令字符串，为命令加上或替换''号，这样操作更安全，不会在解析的时候意外被展开或引用
        safe_command = shlex.quote(args.c)
        
        # 组合完整命令
        full_command = f"{env_prefix} {timeout_prefix} {sudo_prefix} bash -c {safe_command}"

    elif args.s:
        interpreter = 'bash' if args.s.endswith('.sh') else 'python3'

        with open(args.s, 'r', encoding='utf-8') as f:
            script_content = f.read()
            elog.info(f"读取脚本文件成功")
            # 使用heredoc传递脚本内容
            heredoc_tag = f"EOF_{int(time.time())}"  # 唯一标识符
            full_script_content = f"{interpreter} - <<'{heredoc_tag}'\n{script_content}\n{heredoc_tag}"
            elog.info(f"使用EOF包装脚本内容成功")
        
        # 组合完整命令
        full_command = f"{env_prefix} {timeout_prefix} {sudo_prefix} {full_script_content}"
    
    else:
        elog.error(f"参数出现严重异常，args.c 和 args.s 不能同时为空")
        utils.print_error_informantion_and_exit("add_env_sudo_to_commands","参数出现严重异常，args.c 和 args.s 不能同时为空")

    elog.success(f"拼接完成，完整命令: {full_command}，解释器: {interpreter}")
    return full_command.strip()


def handle_interrupt(signum, frame):
    """全局中断信号处理器"""
    global asyncssh_interrupted
    if not asyncssh_interrupted:
        asyncssh_interrupted = True
        # 关键一步：立即屏蔽后续所有Ctrl+C
        signal.signal(signal.SIGINT, signal.SIG_IGN)
        print(f"{config.COLOR_RED}收到中断信号，正在停止任务...{config.COLOR_RESET}")
        elog.warning("收到中断信号，正在停止任务...(handle_interrupt)")


